{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Solving Knapsack Problem with Amazon SageMaker RL"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Knapsack is a canonical operations research problem. We start with a bag and a set of items. We choose which items to put in the bag. Our objective is to maximize the value of the items in the bag; but we cannot put all the items in as the bag capacity is limited. The problem is hard because the items have different values and weights, and there are many combinations to consider.\n",
    "In the classic version of the problem, we pick the items in one shot. But in this baseline, we instead consider the items one at a time over a fixed time horizon."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Statement\n",
    "\n",
    "We start with an empty bag and an item. We need to either put the item in the bag or throw it away. If we put it in the bag, we get a reward equal to the value of the item. If we throw the item away, we get a fixed penalty. In case the bag is too full to accommodate the item, we are forced to throw it away.\n",
    "In the next step, another item appears and we need to decide again if we want to put it in the bag or throw it away.  This process repeats for a fixed number of steps.\n",
    "Since we do not know the value and weight of items that will come in the future, and the bag can only hold so many items, it is not obvious what is the right thing to do."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "At each time step, our agent is aware of the following information:\n",
    "- Weight capacity of the bag\n",
    "- Volume capacity of the bag\n",
    "- Sum of item weight in the bag\n",
    "- Sum of item volume in the bag\n",
    "- Sum of item value in the bag\n",
    "- Current item weight\n",
    "- Current item volume\n",
    "- Current item value\n",
    "- Time remaining\n",
    "\n",
    "At each time step, our agent can take one of the following actions:\n",
    "- Put the item in the bag\n",
    "- Throw the item away\n",
    "\n",
    "At each time step, our agent gets the following reward depending on their action:\n",
    "- Item value if you put it in the bag and bag does not overflow\n",
    "- A penalty if you throw the item away or if the item does not fit in the bag\n",
    "\n",
    "The time horizon is 20 steps. You can see the specifics in the `KnapSackMediumEnv` class in `knapsack_env.py`. There are a couple of other classes that provide an easier (`KnapSackEnv`) and a more difficult version (`KnapSackHardEnv`) of this problem. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Amazon SageMaker RL\n",
    "\n",
    "Amazon SageMaker RL allows you to train your RL agents in cloud machines using docker containers. You do not have to worry about setting up your machines with the RL toolkits and deep learning frameworks. You can easily switch between many different machines setup for you, including powerful GPU machines that give a big speedup. You can also choose to use multiple machines in a cluster to further speedup training, often necessary for production level loads.\n",
    "\n",
    "### Pre-requsites\n",
    "#### Imports\n",
    "\n",
    "To get started, we'll import the Python libraries we need, set up the environment with a few prerequisites for permissions and configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "import boto3\n",
    "import sys\n",
    "import os\n",
    "import glob\n",
    "import re\n",
    "import subprocess\n",
    "from IPython.display import HTML\n",
    "import time\n",
    "from time import gmtime, strftime\n",
    "sys.path.append(\"common\")\n",
    "from misc import get_execution_role, wait_for_s3_object\n",
    "from sagemaker.rl import RLEstimator, RLToolkit, RLFramework"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Settings\n",
    "\n",
    "You can run this notebook from your local host or from a SageMaker notebook instance. In both of these scenarios, you can run the following in either `local` or `SageMaker` modes. The `local` mode uses the SageMaker Python SDK to run your code in a local container before deploying to SageMaker. This can speed up iterative testing and debugging while using the same familiar Python SDK interface. You just need to set `local_mode = True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# run in local mode?\n",
    "local_mode = False\n",
    "\n",
    "# create unique job name \n",
    "job_name_prefix = 'rl-knapsack'\n",
    "\n",
    "# S3 bucket\n",
    "sage_session = sagemaker.session.Session()\n",
    "s3_bucket = sage_session.default_bucket()  \n",
    "print(\"Using s3 bucket %s\" % s3_bucket)  # create this bucket if it doesn't exist\n",
    "s3_output_path = 's3://{}/'.format(s3_bucket) # SDK appends the job name and output folder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Install docker for `local` mode\n",
    "\n",
    "In order to work in `local` mode, you need to have docker installed. When running from you local instance, please make sure that you have docker or docker-compose (for local CPU machines) and nvidia-docker (for local GPU machines) installed. Alternatively, when running from a SageMaker notebook instance, you can simply run the following script \n",
    "\n",
    "Note, you can only run a single local notebook at one time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if local_mode:\n",
    "    !/bin/bash ./common/setup.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create an IAM role\n",
    "Either get the execution role when running from a SageMaker notebook `role = sagemaker.get_execution_role()` or, when running locally, set it to an IAM role with `AmazonSageMakerFullAccess` and `CloudWatchFullAccess permissions`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    role = sagemaker.get_execution_role()\n",
    "except:\n",
    "    role = get_execution_role()\n",
    "\n",
    "print(\"Using IAM role arn: {}\".format(role))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Setup the environment\n",
    "\n",
    "The environment is defined in a Python file called `knapsack_env.py` in the `./src` directory. It implements the init(), step(), reset() and render() functions that describe how the environment behaves. This is consistent with Open AI Gym interfaces for defining an environment.\n",
    "\n",
    "- Init() - initialize the environment in a pre-defined state\n",
    "- Step() - take an action on the environment\n",
    "- reset()- restart the environment on a new episode\n",
    "- render() - get a rendered image of the environment in its current state\n",
    "\n",
    "#### Configure the presets for RL algorithm\n",
    "\n",
    "The presets that configure the RL training jobs are defined in the `preset-knapsack-clippedppo.py` in the `./src` directory. Using the preset file, you can define agent parameters to select the specific agent algorithm. You can also set the environment parameters, define the schedule and visualization parameters, and define the graph manager. The schedule presets will define the number of heat up steps, periodic evaluation steps, training steps between evaluations.\n",
    "\n",
    "These can be overridden at runtime by specifying the RLCOACH_PRESET hyperparameter. Additionally, it can be used to define custom hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize src/preset-knapsack-clippedppo.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Write the Training Code\n",
    "\n",
    "The training code is in the file `train-coach.py` which is also the `./src` directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize src/train-coach.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train the model using Python SDK/ script mode\n",
    "\n",
    "If you are using local mode, the training will run on the notebook instance. When using SageMaker for training, you can select a GPU or CPU instance. The RLEstimator is used for training RL jobs.\n",
    "\n",
    "- Specify the source directory where the environment, presets and training code is uploaded.\n",
    "- Specify the entry point as the training code\n",
    "- Specify the choice of RL toolkit and framework. This automatically resolves to the ECR path for the RL Container.\n",
    "- Define the training parameters such as the instance count, job name, S3 path for output and job name.\n",
    "- Specify the hyperparameters for the RL agent algorithm. The RLCOACH_PRESET can be used to specify the RL agent algorithm you want to use.\n",
    "- Define the metrics definitions that you are interested in capturing in your logs. These can also be visualized in CloudWatch and SageMaker Notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "if local_mode:\n",
    "    instance_type = 'local'\n",
    "else:\n",
    "    instance_type = \"ml.m4.4xlarge\"\n",
    "    \n",
    "estimator = RLEstimator(entry_point=\"train-coach.py\",\n",
    "                        source_dir='src',\n",
    "                        dependencies=[\"common/sagemaker_rl\"],\n",
    "                        toolkit=RLToolkit.COACH,\n",
    "                        toolkit_version='0.11.0',\n",
    "                        framework=RLFramework.TENSORFLOW,\n",
    "                        role=role,\n",
    "                        train_instance_type=instance_type,\n",
    "                        train_instance_count=1,\n",
    "                        output_path=s3_output_path,\n",
    "                        base_job_name=job_name_prefix,\n",
    "                        hyperparameters = {\n",
    "                          \"RLCOACH_PRESET\":\"preset-knapsack-clippedppo\",\n",
    "                          \"rl.agent_params.algorithm.discount\": 0.9,\n",
    "                          \"rl.evaluation_steps:EnvironmentEpisodes\": 8,\n",
    "                        }\n",
    "                    )\n",
    "\n",
    "estimator.fit(wait=local_mode)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Store intermediate training output and model checkpoints\n",
    "\n",
    "The output from the training job above is stored on S3. The intermediate folder contains gifs and metadata of the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name=estimator._current_job_name\n",
    "print(\"Job name: {}\".format(job_name))\n",
    "\n",
    "s3_url = \"s3://{}/{}\".format(s3_bucket,job_name)\n",
    "\n",
    "if local_mode:\n",
    "    output_tar_key = \"{}/output.tar.gz\".format(job_name)\n",
    "else:\n",
    "    output_tar_key = \"{}/output/output.tar.gz\".format(job_name)\n",
    "\n",
    "intermediate_folder_key = \"{}/output/intermediate\".format(job_name)\n",
    "output_url = \"s3://{}/{}\".format(s3_bucket, output_tar_key)\n",
    "intermediate_url = \"s3://{}/{}\".format(s3_bucket, intermediate_folder_key)\n",
    "\n",
    "print(\"S3 job path: {}\".format(s3_url))\n",
    "print(\"Output.tar.gz location: {}\".format(output_url))\n",
    "print(\"Intermediate folder path: {}\".format(intermediate_url))\n",
    "    \n",
    "tmp_dir = \"/tmp/{}\".format(job_name)\n",
    "os.system(\"mkdir {}\".format(tmp_dir))\n",
    "print(\"Create local folder {}\".format(tmp_dir))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualization\n",
    "#### Plot metrics for training job\n",
    "\n",
    "We can pull the reward metric of the training and plot it to see the performance of the model over time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import pandas as pd\n",
    "\n",
    "csv_file_name = \"worker_0.simple_rl_graph.main_level.main_level.agent_0.csv\"\n",
    "key = intermediate_folder_key + \"/\" + csv_file_name\n",
    "wait_for_s3_object(s3_bucket, key, tmp_dir)\n",
    "\n",
    "csv_file = \"{}/{}\".format(tmp_dir, csv_file_name)\n",
    "df = pd.read_csv(csv_file)\n",
    "df = df.dropna(subset=['Training Reward'])\n",
    "x_axis = 'Episode #'\n",
    "y_axis = 'Training Reward'\n",
    "\n",
    "plt = df.plot(x=x_axis,y=y_axis, figsize=(12,5), legend=True, style='b-')\n",
    "plt.set_ylabel(y_axis);\n",
    "plt.set_xlabel(x_axis);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualize the rendered gifs\n",
    "The latest gif file found in the gifs directory is displayed. You can replace the tmp.gif file below to visualize other files generated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = intermediate_folder_key + '/gifs'\n",
    "wait_for_s3_object(s3_bucket, key, tmp_dir)    \n",
    "print(\"Copied gifs files to {}\".format(tmp_dir))\n",
    "\n",
    "glob_pattern = os.path.join(\"{}/*.gif\".format(tmp_dir))\n",
    "gifs = [file for file in glob.iglob(glob_pattern, recursive=True)]\n",
    "extract_episode = lambda string: int(re.search('.*episode-(\\d*)_.*', string, re.IGNORECASE).group(1))\n",
    "gifs.sort(key=extract_episode)\n",
    "print(\"GIFs found:\\n{}\".format(\"\\n\".join([os.path.basename(gif) for gif in gifs])))    \n",
    "\n",
    "# visualize a specific episode\n",
    "gif_index = -1 # since we want last gif\n",
    "gif_filepath = gifs[gif_index]\n",
    "gif_filename = os.path.basename(gif_filepath)\n",
    "print(\"Selected GIF: {}\".format(gif_filename))\n",
    "os.system(\"mkdir -p ./src/tmp_render/ && cp {} ./src/tmp_render/{}.gif\".format(gif_filepath, gif_filename))\n",
    "HTML('<img src=\"./src/tmp_render/{}.gif\">'.format(gif_filename))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluation of RL models\n",
    "\n",
    "We use the last checkpointed model to run evaluation for the RL Agent. \n",
    "\n",
    "#### Load checkpointed model\n",
    "\n",
    "Checkpointed data from the previously trained models will be passed on for evaluation / inference in the checkpoint channel. In local mode, we can simply use the local directory, whereas in the SageMaker mode, it needs to be moved to S3 first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wait_for_s3_object(s3_bucket, output_tar_key, tmp_dir)  \n",
    "\n",
    "if not os.path.isfile(\"{}/output.tar.gz\".format(tmp_dir)):\n",
    "    raise FileNotFoundError(\"File output.tar.gz not found\")\n",
    "os.system(\"tar -xvzf {}/output.tar.gz -C {}\".format(tmp_dir, tmp_dir))\n",
    "\n",
    "if local_mode:\n",
    "    checkpoint_dir = \"{}/data/checkpoint\".format(tmp_dir)\n",
    "else:\n",
    "    checkpoint_dir = \"{}/checkpoint\".format(tmp_dir)\n",
    "\n",
    "print(\"Checkpoint directory {}\".format(checkpoint_dir))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if local_mode:\n",
    "    checkpoint_path = 'file://{}'.format(checkpoint_dir)\n",
    "    print(\"Local checkpoint file path: {}\".format(checkpoint_path))\n",
    "else:\n",
    "    checkpoint_path = \"s3://{}/{}/checkpoint/\".format(s3_bucket, job_name)\n",
    "    if not os.listdir(checkpoint_dir):\n",
    "        raise FileNotFoundError(\"Checkpoint files not found under the path\")\n",
    "    os.system(\"aws s3 cp --recursive {} {}\".format(checkpoint_dir, checkpoint_path))\n",
    "    print(\"S3 checkpoint file path: {}\".format(checkpoint_path))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Run the evaluation step\n",
    "\n",
    "Use the checkpointed model to run the evaluation step. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator_eval = RLEstimator(role=role,\n",
    "                      source_dir='src/',\n",
    "                      dependencies=[\"common/sagemaker_rl\"],\n",
    "                      toolkit=RLToolkit.COACH,\n",
    "                      toolkit_version='0.11.0',\n",
    "                      framework=RLFramework.TENSORFLOW,\n",
    "                      entry_point=\"evaluate-coach.py\",\n",
    "                      train_instance_count=1,\n",
    "                      train_instance_type=instance_type,\n",
    "                      hyperparameters = {\n",
    "                                 \"RLCOACH_PRESET\":\"preset-knapsack-clippedppo\",\n",
    "                                 \"evaluate_steps\": 250, #5 episodes\n",
    "                             }\n",
    "                    )\n",
    "estimator_eval.fit({'checkpoint': checkpoint_path})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the output \n",
    "\n",
    "Optionally, you can run the steps defined earlier to visualize the output "
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "kernelspec": {
   "display_name": "conda_tensorflow_p36",
   "language": "python",
   "name": "conda_tensorflow_p36"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  },
  "notice": "Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved. Licensed under the Apache License, Version 2.0 (the \"License\"). You may not use this file except in compliance with the License. A copy of the License is located at http://aws.amazon.com/apache2.0/ or in the \"license\" file accompanying this file. This file is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
